package git.soulbgm;

import io.netty.bootstrap.Bootstrap;
import io.netty.channel.*;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.bytes.ByteArrayEncoder;
import org.pcap4j.core.*;
import org.pcap4j.packet.*;

import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.*;
import java.util.concurrent.*;

public class Main {
    private static final String TARGET_IP = "127.0.0.1";
    private static final int TARGET_PORT = 8080;
    private static final String FORWARD_IP = "127.0.0.1";
    private static final int FORWARD_PORT = 9090;
    private static final String BIN_HOST = "127.0.0.1";

    // 超时时间，单位毫秒
    private static final long TIMEOUT_MS = 5000;
    // 2^31
    private static final long HALF_MAX_SEQ = 0x80000000L;

    private static Channel forwardChannel;
    // 用于存储 TCP 数据包的序列号和负载数据
    private static final TreeMap<Long, byte[]> packetBuffer = new TreeMap<>();
    // 记录下一个期望的序列号
    private static long nextExpectedSeq = 0;
    // 用于记录数据包的接收时间
    private static final Map<Long, Long> packetReceiveTime = new HashMap<>();
    // 定时任务线程池，用于处理超时的数据包
    private static final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);
    private static PcapHandle handle;

    public static void main(String[] args) {
        try {
            // 初始化 Netty 客户端
            initNettyClient();

            PcapNetworkInterface nif = Pcaps.getDevByAddress(InetAddress.getByName(BIN_HOST));
            System.out.println("使用网络接口: " + nif.getDescription());

            handle = nif.openLive(65536, PcapNetworkInterface.PromiscuousMode.PROMISCUOUS, 10);
            String filter = "tcp and host " + TARGET_IP + " and port " + TARGET_PORT;
            handle.setFilter(filter, BpfProgram.BpfCompileMode.OPTIMIZE);

            // 启动定时任务检查超时数据包
            scheduler.scheduleAtFixedRate(Main::checkTimeoutPackets, TIMEOUT_MS, TIMEOUT_MS, TimeUnit.MILLISECONDS);

            // 开始抓包
            PacketListener listener = packet -> {
                if (packet.contains(TcpPacket.class)) {
                    TcpPacket tcpPacket = packet.get(TcpPacket.class);
                    if (tcpPacket.getPayload() != null) {
                        byte[] payload = tcpPacket.getPayload().getRawData();
                        long seq = tcpPacket.getHeader().getSequenceNumberAsLong();
                        // 处理重复数据包
                        if (packetBuffer.containsKey(seq)) {
                            return;
                        }
                        // 将数据包添加到缓冲区
                        packetBuffer.put(seq, payload);
                        // 记录接收时间
                        packetReceiveTime.put(seq, System.currentTimeMillis());
                        // 尝试重组并转发数据
                        processPacketBuffer();
                    }
                }
            };

            handle.loop(-1, listener);
        } catch (PcapNativeException | NotOpenException e) {
            System.err.println("初始化抓包器或连接 Netty 客户端时出错: " + e.getMessage());
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            System.err.println("抓包过程被中断: " + e.getMessage());
        } catch (UnknownHostException e) {
            throw new RuntimeException(e);
        } finally {
            shutdown();
        }
    }

    private static boolean isSeqAfter(long seq1, long seq2) {
        return (seq1 > seq2 && seq1 - seq2 < HALF_MAX_SEQ) || (seq1 < seq2 && seq2 - seq1 > HALF_MAX_SEQ);
    }

    private static void processPacketBuffer() {
        while (!packetBuffer.isEmpty() && !isSeqAfter(nextExpectedSeq, packetBuffer.firstKey())) {
            Map.Entry<Long, byte[]> entry = packetBuffer.pollFirstEntry();
            long seq = entry.getKey();
            byte[] payload = entry.getValue();
            packetReceiveTime.remove(seq);
            if (forwardChannel != null && forwardChannel.isActive()) {
                forwardChannel.writeAndFlush(payload);
            }
            nextExpectedSeq = (nextExpectedSeq + payload.length) & 0xFFFFFFFFL;
        }
    }

    private static void checkTimeoutPackets() {
        long currentTime = System.currentTimeMillis();
        Iterator<Map.Entry<Long, Long>> iterator = packetReceiveTime.entrySet().iterator();
        while (iterator.hasNext()) {
            Map.Entry<Long, Long> entry = iterator.next();
            long seq = entry.getKey();
            long receiveTime = entry.getValue();
            if (currentTime - receiveTime > TIMEOUT_MS) {
                if (isSeqAfter(nextExpectedSeq, seq)) {
                    // 已经过时的数据包，直接移除
                    packetBuffer.remove(seq);
                    iterator.remove();
                } else {
                    // 丢失的数据包，记录日志，实际可实现重传逻辑
                    System.err.println("Packet with sequence number " + seq + " has timed out.");
                }
            }
        }
    }

    private static void initNettyClient() {
        EventLoopGroup group = new NioEventLoopGroup();
        try {
            Bootstrap b = new Bootstrap();
            b.group(group)
                    .channel(NioSocketChannel.class)
                    .option(ChannelOption.SO_KEEPALIVE, true)
                    .handler(new ChannelInitializer<SocketChannel>() {
                        @Override
                        public void initChannel(SocketChannel ch) {
                            ch.pipeline().addLast(new ByteArrayEncoder());
                        }
                    });

            // 连接到目标服务器
            ChannelFuture f = b.connect(FORWARD_IP, FORWARD_PORT).sync();
            forwardChannel = f.channel();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static void shutdown() {
        if (handle != null && handle.isOpen()) {
            handle.close();
        }
        scheduler.shutdown();
        try {
            if (!scheduler.awaitTermination(5, TimeUnit.SECONDS)) {
                scheduler.shutdownNow();
            }
        } catch (InterruptedException e) {
            scheduler.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }
}